import wx
import copy
from Buffers import buffers

# distance_transform_se = ((1, 1, 1), (1, 1, 1), (1, 1, 1))

class MorphologyFrame(wx.Frame):
	def __init__(self, parent, id, title, pos):
		wx.Frame.__init__(self, parent, id, title, pos, style=wx.FRAME_NO_TASKBAR)
		panel = wx.Panel(self, wx.ID_ANY)
		self.panel = panel
		self.vals_ctrl = []

		spinner_width = wx.SpinCtrl(panel, wx.ID_ANY, "", (0, 0), (50, -1))
		spinner_width.SetRange(1, 20)
		spinner_width.SetValue(5)
		spinner_height = wx.SpinCtrl(panel, wx.ID_ANY, "", (0, 0), (50, -1))
		spinner_height.SetRange(1, 20)
		spinner_height.SetValue(5)
		spinner_center_x = wx.SpinCtrl(panel, wx.ID_ANY, "", (0, 0), (50, -1))
		spinner_center_x.SetRange(1, 20)
		spinner_center_x.SetValue(2)
		spinner_center_y = wx.SpinCtrl(panel, wx.ID_ANY, "", (0, 0), (50, -1))
		spinner_center_y.SetRange(1, 20)
		spinner_center_y.SetValue(2)
		self.Bind(wx.EVT_TEXT, self.OnChangeWidth, spinner_width)
		self.Bind(wx.EVT_TEXT, self.OnChangeHeight, spinner_height)
		self.Bind(wx.EVT_TEXT, self.OnChangeCenter_x, spinner_width)
		self.Bind(wx.EVT_TEXT, self.OnChangeCenter_y, spinner_height)

		self.filter_dim = [spinner_width.GetValue(), spinner_height.GetValue()]
		self.filter_center = [spinner_center_x.GetValue(), spinner_center_y.GetValue()]

		self.sizer_grid = self.create_grid()

		button_dilation = wx.Button(panel, wx.ID_ANY, "dilation", style=wx.BU_EXACTFIT)
		button_erosion = wx.Button(panel, wx.ID_ANY, "erosion", style=wx.BU_EXACTFIT)
		button_opening = wx.Button(panel, wx.ID_ANY, "opening", style=wx.BU_EXACTFIT)
		button_closing = wx.Button(panel, wx.ID_ANY, "closing", style=wx.BU_EXACTFIT)
		button_distance_transform = wx.Button(panel, wx.ID_ANY, "distance_transform", style=wx.BU_EXACTFIT)
		button_skeleton = wx.Button(panel, wx.ID_ANY, "skeleton", style=wx.BU_EXACTFIT)
		button_reconstruction = wx.Button(panel, wx.ID_ANY, "reconstruction", style=wx.BU_EXACTFIT)
		self.Bind(wx.EVT_BUTTON, self.OnDilation, button_dilation)
		self.Bind(wx.EVT_BUTTON, self.OnErosion, button_erosion)
		self.Bind(wx.EVT_BUTTON, self.OnOpening, button_opening)
		self.Bind(wx.EVT_BUTTON, self.OnClosing, button_closing)
		self.Bind(wx.EVT_BUTTON, self.OnDistanceTransform, button_distance_transform)
		self.Bind(wx.EVT_BUTTON, self.OnSkeleton, button_skeleton)
		self.Bind(wx.EVT_BUTTON, self.OnReconstruction, button_reconstruction)

		button_cancel = wx.Button(panel, wx.ID_ANY, "Cancel", style=wx.BU_EXACTFIT)
		self.Bind(wx.EVT_BUTTON, self.OnCancel, button_cancel)

		sizer_spinner = wx.BoxSizer(wx.HORIZONTAL)
		sizer_spinner.Add(spinner_width)
		sizer_spinner.AddSpacer(1)
		sizer_spinner.Add(spinner_height)
		sizer_spinner.AddSpacer(5)
		sizer_spinner.Add(spinner_center_x)
		sizer_spinner.AddSpacer(1)
		sizer_spinner.Add(spinner_center_y)

		sizer_top_left = wx.BoxSizer(wx.VERTICAL)
		sizer_top_left.Add(sizer_spinner)
		sizer_top_left.AddSpacer(20)
		sizer_top_left.Add(self.sizer_grid, flag=wx.ALIGN_CENTER_HORIZONTAL)

		sizer_top_right = wx.BoxSizer(wx.VERTICAL)
		sizer_top_right.Add(button_dilation)
		sizer_top_right.Add(button_erosion)
		sizer_top_right.Add(button_opening)
		sizer_top_right.Add(button_closing)
		sizer_top_right.Add(button_distance_transform)
		sizer_top_right.Add(button_skeleton)
		sizer_top_right.Add(button_reconstruction)

		sizer_top = wx.BoxSizer(wx.HORIZONTAL)
		sizer_top.AddSpacer(20)
		sizer_top.Add(sizer_top_left)
		sizer_top.AddSpacer(20)
		sizer_top.Add(sizer_top_right)

		sizer_bottom = wx.BoxSizer(wx.HORIZONTAL)
		sizer_bottom.Add(button_cancel)

		self.sizer = wx.BoxSizer(wx.VERTICAL)
		self.sizer.Add(sizer_top)
		self.sizer.Add(sizer_bottom, flag=wx.ALIGN_RIGHT)
		self.sizer.Fit(self)

		panel.SetSizer(self.sizer)
		self.Bind(wx.EVT_CLOSE, self.OnCloseWindow)

	def set_source(self, source_name, source_image):
		self.source_name = source_name
		self.source_image = source_image.ConvertToGreyscale()

	def OnCancel(self, event):
		self.Hide()
	
	def OnCloseWindow(self, event):
		self.Hide()

	def OnChangeWidth(self, event):
		self.filteridth = event.GetInt()
		self.change_shape()
		
	def OnChangeHeight(self, event):
		self.filter_dim[1] = event.GetInt()
		self.change_shape()

	def create_grid(self):
		self.vals_ctrl = []
		for i in range(self.filter_dim[0]*self.filter_dim[1]):
			self.vals_ctrl.append(wx.TextCtrl(self.panel, wx.ID_ANY, value="1", size=(40, -1)))
		
		sizer_grid = wx.GridSizer(rows=self.filter_dim[1], cols=self.filter_dim[0], hgap=2, vgap=2)
		for val_ctrl in self.vals_ctrl:
			sizer_grid.Add(val_ctrl)

		return sizer_grid

	def change_shape(self):
		self.sizer_grid = self.create_grid()
		self.sizer.Fit(self)

	def OnChangeCenter_x(self, event):
		self.filter_center[0] = event.GetInt()

	def OnChangeCenter_y(self, event):
		self.filter_center[1] = event.GetInt()

	def get_filter(self):
		filter = []
		for i in range(self.filter_dim[1]):
			row = []
			base = i*self.filter_dim[0]
			for j in range(self.filter_dim[0]):
				idx = base + j
				val = eval(self.vals_ctrl[idx].GetValue())
				row.append(val)
			filter.append(row)
		return filter

	def display(self, image, postfix):
		name = self.addto_buffer(image, postfix)
		self.GetParent().displayBuffer(name)
		self.Hide()

	def addto_buffer(self, image, postfix):
		new_name = self.source_name + " - " + postfix
		name = buffers.add(new_name, image, self.source_name)
		return name

	def OnDilation(self, event):
		filter = self.get_filter()
		image = self.dilation(self.source_image, filter)
		self.display(image, "dilation")

	def dilation(self, source, filter):
		gray_source = source
		source_w, source_h  = (gray_source.GetWidth(), gray_source.GetHeight())
		w = source_w + self.filter_dim[0] - 1
		h = source_h + self.filter_dim[1] - 1

		image = wx.EmptyImage(w, h)
		for x in range(w):
			for y in range(h):
				max = 0
				for i in range(self.filter_dim[0]):
					for j in range(self.filter_dim[1]):
						diff_x = i - self.filter_center[0]
						diff_y = j - self.filter_center[1]
						filter_val = filter[i][j]

						if (x-i in range(source_w)) and (y-j in range(source_h)):
							val = filter_val + gray_source.GetRed(x-i, y-j)
						else:
							val = 0

						if val > max:
							max = val

				if max > 255:
					max = 255 
				elif max < 0:
					max = 0
				val = int(max)
				image.SetRGB(x, y, val, val, val)
		return image

	def OnErosion(self, event):
		filter = self.get_filter()
		image = self.erosion(self.source_image, filter)
		self.display(image, "erosion")

	def erosion(self, source, filter):
		gray_source = source
		source_w, source_h  = (gray_source.GetWidth(), gray_source.GetHeight())
		w = source_w - self.filter_dim[0] + 1
		h = source_h - self.filter_dim[1] + 1

		image = wx.EmptyImage(w, h)
		for x in range(w):
			for y in range(h):
				min = 255
				for i in range(self.filter_dim[0]):
					for j in range(self.filter_dim[1]):
						diff_x = i - self.filter_center[0]
						diff_y = j - self.filter_center[1]
						filter_val = filter[i][j]

						if (x+i in range(source_w)) and (y+j in range(source_h)):
							val = filter_val + gray_source.GetRed(x+i, y+j)
						else:
							val = 0

						if val < min:
							min = val

				if min > 255:
					min = 255 
				elif min < 0:
					min = 0
				val = int(min)
				image.SetRGB(x, y, val, val, val)
		return image

	def OnOpening(self, event):
		filter = self.get_filter()
		tmp = self.erosion(self.source_image, filter)
		image = self.dilation(tmp, filter)
		self.display(image, "opening")

	def OnClosing(self, event):
		filter = self.get_filter()
		tmp = self.dilation(self.source_image, filter)
		image = self.erosion(tmp, filter)
		self.display(image, "closing")

	def OnDistanceTransform(self, event):
		image = self.distance_transform(self.source_image)
		self.display(image, "distance_transform")

	def OnSkeleton(self, event):
		dis_name = self.source_name + " - distance_transform"
		if not buffers.has_name(dis_name):
			dis_image = self.distance_transform(self.source_image)
			self.addto_buffer(dis_image, "distance_transform")
		else:
			dis_image = buffers.get_image_by_name(dis_name)
		image = self.skeleton(dis_image)
		self.display(image, "skeleton")
	
	def OnReconstruction(self, event):
		if not self.source_name.endswith('skeleton'):
			return
		image = self.reconstruction(self.source_image)
		self.display(image, "reconstruction")
	
	def distance_transform(self, source):
		self.filter_dim = (3, 3)
		mono_source = source.ConvertToMono(255, 255, 255)
		#mono_source = source.ConvertToMono(0, 0, 0)

		w, h = (mono_source.GetWidth(), mono_source.GetHeight())
		image = wx.EmptyImage(w, h)

		count = 0
		change_flag = True
		tmp = mono_source
		gray_step = 10
		while(change_flag):
			change_flag = False
			for x in range(w):
				for y in range(h):
					gray = tmp.GetRed(x, y)
					if (gray == count) and \
							(len(self.neighbor_pos(tmp, x, y)) == 9) and \
							self.neighbor_same(tmp, x, y):
						new_gray = count + gray_step
						if new_gray > 255:
							new_gray = 255
						image.SetRGB(x, y, new_gray, new_gray, new_gray)
						change_flag = True
			count += gray_step
			tmp = image.Copy()
		return image

	def skeleton(self, dis_image):
		w, h = (dis_image.GetWidth(), dis_image.GetHeight())
		image = wx.EmptyImage(w, h)

		for x in range(w):
			for y in range(h):
				val = dis_image.GetRed(x, y)
				pos = self.neighbor_pos(dis_image, x, y)
				for tup in pos:
					if val < dis_image.GetRed(*tup):
						break
				else:
					image.SetRGB(x, y, val, val, val)

		return image

	def reconstruction(self, skeleton_image):
		w, h = (skeleton_image.GetWidth(), skeleton_image.GetHeight())
		image = wx.EmptyImage(w, h)

		gray_step = 10
		stack = []
		for x in range(w):
			for y in range(h):
				gray = skeleton_image.GetRed(x, y)
				if gray > gray_step:
					stack.append(((x, y), gray))
					
		while(len(stack) > 0):
			tup = stack.pop()
			image.SetRGB(tup[0][0], tup[0][1], tup[1], tup[1], tup[1])
			if tup[1] > gray_step:
				neighbor = self.neighbor_pos(image, *tup[0])
				for pos in neighbor:
					if tup[1]-gray_step > image.GetRed(*pos):
						stack.append((pos, tup[1]-gray_step))

		return image

	def is_empty(self, image):
		w, h = (image.GetWidth(), image.GetHeight())
		for i in range(w):
			for j in range(h):
				if (image.GetRed(i, j) != 0)\
						or (image.GetGreen(i, j) != 0)\
						or (image.GetBlue(i, j) != 0):
					return False
		return True

	def neighbor_same(self, mono_image, x, y):
		pos = self.neighbor_pos(mono_image, x, y)
		for tup in pos:
			val = mono_image.GetRed(*tup)
			try:
				if last != val:
					return False
			except:
				pass
			last = val
		return True

	def neighbor_pos(self, mono_image, x, y):
		w, h = (mono_image.GetWidth(), mono_image.GetHeight())
		pos = []

		if 0 <= x-1 < w and 0 <= y-1 < h:
			pos.append((x-1, y-1))
		if 0 <= x-1 < w and 0 <= y < h:
			pos.append((x-1, y))
		if 0 <= x-1 < w and 0 <= y+1 < h:
			pos.append((x-1, y+1))

		if 0 <= x < w and 0 <= y-1 < h:
			pos.append((x, y-1))
		if 0 <= x < w and 0 <= y < h:
			pos.append((x, y))
		if 0 <= x < w and 0 <= y+1 < h:
			pos.append((x, y+1))

		if 0 <= x+1 < w and 0 <= y-1 < h:
			pos.append((x+1, y-1))
		if 0 <= x+1 < w and 0 <= y < h:
			pos.append((x+1, y))
		if 0 <= x+1 < w and 0 <= y+1 < h:
			pos.append((x+1, y+1))

		return pos

	def diff_region(self, image, image_a, image_b, size, count):
		w, h = image_a.GetSize()
		#image = wx.EmptyImage(w, h)
		for i in range(w):
			for j in range(h):
				if (image_a.GetRed(i, j) != image_b.GetRed(i+1, j+1)) \
						or (image_a.GetGreen(i, j) != image_b.GetGreen(i+1, j+1)) \
						or (image_a.GetBlue(i, j) != image_b.GetBlue(i+1, j+1)):
					image.SetRGB(i+count, j+count, count, count, count)

	def fill_color(self, image, region, color):
		w, h = (image.GetWidth(), image.GetHeight())
		for i in range(w):
			for j in range(h):
				if region.GetRed(i, j) != 0:
					image.SetRGB(i, j, *color)
